# 根据参数画线
from math import tan, sin, cos, acos

import matplotlib.pyplot as plt
import numpy as np

from utils.util import get_angle_cos, _k, get_angle_tan

step = 10
red = '#FF0000'  # 直线
blue = '#0000FF'  # 缓和曲线
green = '#00FF00'  # 圆弧


def read_line_data(filename='data/左线平面参数.txt'):
    Xs, Ys, Rs, fcls, pcls = [], [], [], [], []  # x, y, 半径, 前缓和曲线长, 后缓和曲线长
    with open(filename) as data_file:
        for line in data_file:
            if len(line) == 0:
                continue
            y, x, r, fcl, pcl = [float(i) for i in line.replace('\n', '').split('\t')]
            Xs.append(x)
            Ys.append(y)
            Rs.append(r)
            fcls.append(fcl)
            pcls.append(pcl)
    return [Xs, Ys, Rs, fcls, pcls]


def draw_line(data):
    Xs, Ys, Rs, fcls, pcls = data
    x_array, y_array = [], []
    c = []  # 颜色
    x_array_straight, y_array_straight = [Xs[0], Xs[1]], [Ys[0], Ys[1]]
    for i in range(1, len(Xs)-1):
        R = Rs[i]
        start, JD, end = [[Xs[j], Ys[j]] for j in range(i-1, i+2)]
        fcl, pcl = fcls[i], pcls[i]
        # straight_x, straight_y, cl_x, cl_y, curve_x, curve_y = [], [], [], [], [], []

        # ls不为0, 计算前缓和曲线
        alpha = np.pi - get_angle_cos(start, JD, end)  # alpha
        flag = judge_direction(start, JD, end)  # 判断左偏还是右偏

        F_ZH_JD = get_angle_tan(JD, start)
        ZH, HY, beta1 = get_cl_point(fcl, alpha, JD, R, F_ZH_JD, flag)  # 前缓和曲线交点
        straight_x, straight_y = draw_straight(start, ZH)
        cl_x, cl_y = draw_cl(ZH, F_ZH_JD, fcl, R, flag)

        # 直线和前缓和曲线
        x_array += straight_x
        y_array += straight_y
        c += [red for i in straight_x]
        x_array += cl_x
        y_array += cl_y
        c += [blue for i in cl_x]

        # 后缓和曲线
        F_HZ_JD = np.pi + get_angle_tan(JD, end)
        HZ, YH, beta2 = get_cl_point(pcl, alpha, JD, R, F_HZ_JD, flag)  # 后缓和曲线交点
        if i == len(Xs)-2:
            straight_x, straight_y = draw_straight(end, HZ)
        cl_x, cl_y = draw_cl(HZ, F_HZ_JD, fcl, R, flag)

        # 为了防止重复连线
        Xs[i] = HZ[0]
        Ys[i] = HZ[1]

        # 圆弧
        ly = R * (alpha - beta1 - beta2)
        F_HY = F_ZH_JD + flag * beta1
        curve_x, curve_y = draw_curve(HY, F_HY, ly, R, flag)
        x_array += curve_x
        y_array += curve_y
        c += [green for i in curve_x]

        # 后缓和曲线和直线
        x_array += cl_x
        y_array += cl_y
        c += [blue for i in cl_x]
        # 是最后一段才连接最后一段直线
        if i == len(Xs) - 2:
            x_array += straight_x
            y_array += straight_y
            c += [red for i in straight_x]

        # 对照的直线
        x_array_straight.append(end[0])
        y_array_straight.append(end[1])

        print('\r{}/{}'.format(i, len(Xs) - 2), end='')

    return [x_array, y_array, x_array_straight, y_array_straight, c]


def judge_direction(start, JD, end):
    k1 = _k(start, JD)
    k2 = _k(end, JD)
    if k1 > k2:
        return -1
    return 1


def get_cl_point(ls, alpha, JD, R, F, flag):
    # 只有圆弧没有缓和曲线的情况
    if ls == 0:
        T = R / tan((np.pi - alpha)/2)
        x = JD[0] - T * cos(F)
        y = JD[1] - T * sin(F)
        return [[x, y], [x, y], 0]
    x0 = _x(ls, ls, R)
    y0 = _y(x0, R, x0)
    beta = _beta(ls, ls, R)
    s1, s2 = y0 / tan(beta), y0 / sin(beta)
    t = R * tan((alpha - 2*beta) / 2)
    AB = t + s2
    B_JD = AB * sin((np.pi + alpha - 2*beta) / 2) / sin((np.pi - alpha) / 2)
    T = x0 + B_JD - s1
    x = JD[0] - T * cos(F)
    y = JD[1] - T * sin(F)
    return [[x, y], transform([x0, y0], [x, y], F, flag), beta]


def draw_cl(Z, F, ls, R, flag):
    if ls == 0:
        return [[], []]
    l = 0
    x0 = _x(ls, ls, R)
    x_array, y_array = [], []
    while l <= ls:
        x = _x(l, ls, R)
        y = _y(x, R, x0)
        xx, yy = transform([x, y], Z, F, flag)
        x_array.append(xx)
        y_array.append(yy)
        l += step
    return [x_array, y_array]


def draw_straight(p1, p2):
    x0 = min(p1[0], p2[0])
    xt = max(p1[0], p2[0])
    k = _k(p1, p2)
    b = p1[1] - k * p1[0]
    x_array, y_array = [], []
    x = x0
    while x <= xt:
        x_array.append(x)
        y_array.append(k * x + b)
        x += step
    return [x_array, y_array]


def draw_curve(Y, F, ly, R, flag):
    x_array, y_array = [], []
    s = 0
    while s <= ly:
        theta = s / R
        p = [R * sin(theta), R * (1 - cos(theta))]
        Xp, Yp = transform(p, Y, F, flag)
        x_array.append(Xp)
        y_array.append(Yp)
        s += step
    return [x_array, y_array]


def _x(l, ls, R):
    return l * (1 - 1/40 * pow(l, 4) / pow(R * ls, 2))


def _y(x, R, X):
    return pow(x, 3) / (6 * R * X)


def _beta(l, ls, R):
    return pow(l, 2) / (2 * R * ls)


def transform(p, Z, F, flag):
    Xp = Z[0] + p[0] * cos(F) - flag * p[1] * sin(F)
    Yp = Z[1] + p[0] * sin(F) + flag * p[1] * cos(F)
    return [Xp, Yp]


if __name__ == '__main__':
    x_array, y_array, x_array_straight, y_array_straight, c = draw_line(read_line_data())
    print('\rdrawing straight lines...', end='')
    plt.plot(np.array(x_array_straight), np.array(y_array_straight), 'y')  # 直接连线
    plt.scatter(x=np.array(x_array_straight), y=np.array(y_array_straight), s=75, alpha=0.5)  # 点
    print('\rdrawing curve lines...', end='')
    plt.scatter(x=np.array(x_array), y=np.array(y_array), c=c, s=10, alpha=0.5)  # 带曲线的连线
    plt.show()
